import cv2
import layoutparser as lp
import logging


def slice_table(image):
    '''版面分析，将图片分析为表格，并且将图片保存
    img is cv2 image object
    '''
    try:
        image = image[..., ::-1]
        # 加载只版面分析表格
        model = lp.PaddleDetectionLayoutModel(config_path="lp://TableBank/ppyolov2_r50vd_dcn_365e_tableBank_latex/config",
            threshold=0.6,
            label_map={0:"Table"},
            enforce_cpu=True,
            enable_mkldnn=True
        )
        # 检测
        layout = model.detect(image)
        # show img
        # show_img = lp.draw_box(image, layout, box_width=1, show_element_type=True)
        # show_img.show()
        max_x = image.shape[1]
        max_y = image.shape[0]
        if layout.to_dict().get('blocks'):
            img_px = layout.to_dict().get('blocks')[0]
            y1 = int(img_px['y_1']) - 3 if int(img_px['y_1']) - 6 > 0 else 0
            y2 = int(img_px['y_2']) + 3 if int(img_px['y_2']) + 3 < max_y else max_y
            x1 = int(img_px['x_1']) - 3 if int(img_px['x_1']) - 3 > 0 else 0
            x2 = int(img_px['x_2']) + 3 if int(img_px['x_2']) + 3 < max_x else max_x
            table_img = image[y1:y2, x1:x2]
            # cv2.imwrite('108.jpg', new_img)
            # cv2.imshow('table', table_img)
            # cv2.waitKey(0)
            # 截取票号区域
            if int(img_px['y_1']) > 45:
                label_img = image[0:int(img_px['y_1']), x1:x2]
                # cv2.imshow('lable', label_img)
                return table_img, label_img, 'ok'
            # cv2.waitKey(0)
            return table_img, None, 'ok'
        return None, None, 'no data return'
    except Exception as e:
        logging.error(e)
        return None, None, str(e)


if __name__ == "__main__":
    # main(r'D:\aeas\table-detect\new_img\110.jpg')
    import sys
    import os
    current_dir = os.getcwd()
    small_imgs_dir = os.path.join(current_dir, 'label-imgs')
    if not os.path.exists(small_imgs_dir):
        os.makedirs(small_imgs_dir, exist_ok=True)
        os.chmod(small_imgs_dir, mode=0o755)
    if len(sys.argv) > 1:
        arg = sys.argv[1]
        if not os.path.exists(arg) or not os.path.isdir(arg):
            print('输入的路径不存在或不是一个目录')
            sys.exit()
        files = os.listdir(arg)
        for f_path in files:
            try:
                if 'jpg' not in os.path.splitext(f_path)[1]:
                    continue
                image = cv2.imread(os.path.join(arg, f_path))
                table_img, label_img, msg = slice_table(image)
                print(f'{f_path}识别结果：{msg}，是否存在label：{not label_img is None}')
                if not label_img is None:
                    label_path = os.path.join(small_imgs_dir, '%s-label.jpg' % os.path.splitext(f_path)[0])
                    cv2.imwrite(label_path, label_img)
            except Exception as e:
                print(f'处理{f_path}时发生异常：{e}')